import argparse
import os, json
import plotly.graph_objects as go
from tqdm import tqdm

'''
# @Time : 2024/11/23 17:46
# @Author : Shiyu He
# @University : Xinjiang University
'''

parser = argparse.ArgumentParser(description="plot dynamic position")
parser.add_argument("input_path", type=str, help="Path to the input position file.")
args = parser.parse_args()

# 读取 JSON 数据
with open(os.path.join(args.input_path, "state.json"), "r") as file:
    data = json.load(file)

# 提取数据
plane_positions = data["plane_positions"]
enemy_positions = data["enemy_positions"]

# 提取所有 x, y, z 数据并显示进度
plane_x = [pos[0] for pos in plane_positions.values()]
plane_y = [pos[1] for pos in plane_positions.values()]
plane_z = [pos[2] for pos in plane_positions.values()]

enemy_x = [pos[0] for pos in enemy_positions.values()]
enemy_y = [pos[1] for pos in enemy_positions.values()]
enemy_z = [pos[2] for pos in enemy_positions.values()]

# 创建 Plotly 图表
fig = go.Figure()

# 添加飞机的轨迹
fig.add_trace(go.Scatter3d(
    x=[plane_x[0]], y=[plane_y[0]], z=[plane_z[0]],
    mode='lines+markers',
    opacity=0.8,
    line=dict(width=1, color='yellow'),
    marker=dict(size=2, color='red'),
    name='Plane'
))

# 添加敌机的轨迹
fig.add_trace(go.Scatter3d(
    x=[enemy_x[0]], y=[enemy_y[0]], z=[enemy_z[0]],
    mode='lines+markers',
    opacity=0.8,
    line=dict(width=1, color='black'),
    marker=dict(size=2, color='blue'),
    name='Enemy'
))

# 更新布局
fig.update_layout(
    scene=dict(
        xaxis_title='X Position',
        yaxis_title='Y Position',
        zaxis_title='Z Position'
    ),
    title="3D Motion of Plane and Enemy",
    margin=dict(l=0, r=0, b=0, t=40)
)

# 创建动画帧
frames = []
for i in tqdm(range(len(plane_x)), desc="Create Frame:"):
    frame_data = [
        go.Scatter3d(
            x=plane_x[:i+1],
            y=plane_y[:i+1],
            z=plane_z[:i+1],
            mode='lines+markers',
            opacity=0.8,
            line=dict(width=1, color='yellow'),
            marker=dict(size=2, color='red'),
            name='Plane'
        ),
        go.Scatter3d(
            x=enemy_x[:i+1],
            y=enemy_y[:i+1],
            z=enemy_z[:i+1],
            mode='lines+markers',
            opacity=0.8,
            line=dict(width=1, color='black'),
            marker=dict(size=2, color='blue'),
            name='Enemy'
        )
    ]
    frames.append(go.Frame(data=frame_data, name=f"frame{i}"))


# 添加帧到图表
fig.frames = frames

# 动画按钮设置
fig.update_layout(
    updatemenus=[dict(
        type="buttons",
        showactive=False,
        buttons=[
            dict(label="Play",
                 method="animate",
                 args=[None, dict(frame=dict(duration=50, redraw=True), fromcurrent=True)]),
            dict(label="Pause",
                 method="animate",
                 args=[[None], dict(frame=dict(duration=0, redraw=False), mode="immediate")])
        ]
    )]
)

# 添加帧到图表并保存
fig.frames = frames

print("Waiting a few...")
# 显示图表
fig.show()
fig.write_html( "dynamic_plot.html")
print("Animation saved to dynamic_plot.html")
